# 导入异步IO库
import asyncio

# 导入日期时间模块
from datetime import datetime

# 导入类型注解相关
from typing import List, Optional, Any, Literal

# 导入pydantic的数据模型基类和字段定义
from pydantic import BaseModel, Field

# MCP 相关导入
from mcp.server.fastmcp import FastMCP
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client


# 1. MCP 响应结构定义
# 定义内容块基类，继承自pydantic的BaseModel
class ContentBlock(BaseModel):
    """内容块基类"""

    # 占位，不添加实际字段
    pass


# 定义文本内容块，继承自ContentBlock
class TextContent(ContentBlock):
    """文本内容块"""

    # 类型字段，固定为"text"
    type: Literal["text"] = "text"
    # 文本内容
    text: str


# 定义数据内容块，继承自ContentBlock
class DataContent(ContentBlock):
    """数据内容块"""

    # 类型字段，固定为"data"
    type: Literal["data"] = "data"
    # 任意类型的数据
    data: Any
    # 数据的MIME类型
    mimeType: str


# 定义MCP响应结构
class MCPResponse(BaseModel):
    """MCP 响应结构"""

    # 可读内容块列表
    content: List[ContentBlock]  # 可读内容
    # 可选的结构化内容
    structuredContent: Optional[Any]  # 结构化内容（可选）


# 2. 工具返回的结构化内容模型
# 计算表达式结果的数据模型
class CalculationResult(BaseModel):
    """计算表达式结果"""

    # 数学表达式
    expression: str = Field(..., description="数学表达式")
    # 计算结果
    result: float = Field(..., description="计算结果")
    # 计算步骤列表
    steps: List[str] = Field(..., description="计算步骤")
    # 计算时间
    timestamp: datetime = Field(..., description="计算时间")


# 搜索结果的数据模型
class SearchResult(BaseModel):
    """搜索结果"""

    # 结果标题
    title: str = Field(..., description="结果标题")
    # 结果链接
    url: str = Field(..., description="结果链接")
    # 结果摘要
    snippet: str = Field(..., description="结果摘要")
    # 相关度评分，范围0~1
    score: float = Field(..., ge=0, le=1, description="相关度评分")


# 3. 创建 MCP 服务器
# 实例化FastMCP对象，命名为"Structured"
mcp = FastMCP(name="Structured")


# 注册计算表达式工具
@mcp.tool()
async def calculate_expression(expr: str) -> CalculationResult:
    """计算数学表达式并返回结构化结果"""
    try:
        # 使用eval计算表达式
        result = eval(expr)
        # 构建结构化结果并返回
        return CalculationResult(
            expression=expr,
            result=result,
            steps=[f"计算 {expr} = {result}"],
            timestamp=datetime.now(),
        )
    except Exception as e:
        # 捕获异常并抛出ValueError
        raise ValueError(f"计算失败: {e}")


# 注册搜索工具
@mcp.tool()
async def search_web(query: str, limit: int = 5) -> List[SearchResult]:
    """模拟网络搜索并返回结构化结果"""
    # 初始化结果列表
    results = []
    # 生成最多3条模拟搜索结果
    for i in range(min(limit, 3)):
        results.append(
            SearchResult(
                title=f"搜索结果 {i+1}: {query}",
                url=f"https://example.com/result{i+1}",
                snippet=f"这是关于 {query} 的第 {i+1} 个搜索结果...",
                score=0.9 - i * 0.1,
            )
        )
    # 返回搜索结果列表
    return results


# 4. 客户端测试函数
# 测试结构化工具的异步函数
async def test_structured_tools(session: ClientSession):
    """测试结构化工具"""
    # 打印测试计算工具
    print("=== 测试计算工具 ===")
    try:
        # 调用计算表达式工具
        result = await session.call_tool("calculate_expression", {"expr": "2 + 3 * 4"})
        # 打印可读内容
        print("可读内容:")
        for block in result.content:
            # 判断内容块类型并打印
            if isinstance(block, types.TextContent):
                print(f"  文本: {block.text}")
            elif isinstance(block, types.DataContent):
                print(f"  数据: {block.data} (类型: {block.mimeType})")
        # 打印结构化内容
        print("\n结构化内容:")
        if result.structuredContent:
            print(f"  类型: {type(result.structuredContent)}")
            print(f"  内容: {result.structuredContent}")
            # 如果结构化内容为字典，转换为模型并打印详细信息
            if isinstance(result.structuredContent, dict):
                calc_result = CalculationResult(**result.structuredContent)
                print(f"  表达式: {calc_result.expression}")
                print(f"  结果: {calc_result.result}")
                print(f"  步骤: {calc_result.steps}")
                print(f"  时间: {calc_result.timestamp}")
    except Exception as e:
        # 捕获异常并打印错误信息
        print(f"计算工具测试失败: {e}")

    # 打印测试搜索工具
    print("\n=== 测试搜索工具 ===")
    try:
        # 调用搜索工具
        result = await session.call_tool(
            "search_web", {"query": "Python MCP", "limit": 3}
        )
        # 打印可读内容
        print("可读内容:")
        for block in result.content:
            # 判断内容块类型并打印
            if isinstance(block, types.TextContent):
                print(f"  文本: {block.text}")
            elif isinstance(block, types.DataContent):
                print(f"  数据: {block.data} (类型: {block.mimeType})")
        # 打印结构化内容
        print("\n结构化内容:")
        if result.structuredContent:
            print(f"  类型: {type(result.structuredContent)}")
            print(f"  内容: {result.structuredContent}")
            # 如果结构化内容为列表，遍历并打印每个搜索结果
            if isinstance(result.structuredContent, list):
                for i, item in enumerate(result.structuredContent):
                    if isinstance(item, dict):
                        search_result = SearchResult(**item)
                        print(f"  结果 {i+1}:")
                        print(f"    标题: {search_result.title}")
                        print(f"    链接: {search_result.url}")
                        print(f"    摘要: {search_result.snippet}")
                        print(f"    评分: {search_result.score}")
    except Exception as e:
        # 捕获异常并打印错误信息
        print(f"搜索工具测试失败: {e}")


# 5. 客户端主函数
# 客户端运行主函数
async def run_client():
    """运行客户端测试"""
    # 打印启动信息
    print("启动 MCP 客户端测试...")
    # 构造服务器参数，命令为python，参数为当前文件和"serve"
    server_params = StdioServerParameters(command="python", args=[__file__, "serve"])
    # 使用stdio_client连接到服务器
    async with stdio_client(server_params) as (read, write):
        # 创建客户端会话
        async with ClientSession(read, write) as session:
            # 初始化协议
            await session.initialize()
            # 列出可用工具
            tools = await session.list_tools()
            print(f"可用工具: {[t.name for t in tools.tools]}")
            # 测试结构化工具
            await test_structured_tools(session)


# 6. 服务器入口
# 服务器运行入口函数
def run_server():
    """启动 MCP 服务器"""
    # 打印启动信息
    print("启动 MCP 服务器...")
    # 以stdio方式运行MCP服务器
    mcp.run(transport="stdio")


# 7. 主函数
# 程序主入口
def main():
    # 导入sys模块
    import sys

    # 判断命令行参数，决定运行模式
    if len(sys.argv) >= 2 and sys.argv[1] == "serve":
        # 服务器模式
        run_server()
    else:
        # 客户端模式
        print("启动 MCP 客户端测试...")
        print("将自动启动服务器进程并测试结构化工具")
        asyncio.run(run_client())


# 判断是否为主模块，若是则执行main函数
if __name__ == "__main__":
    main()
